package stream

import (
	"sync"

	"github.com/wandb/wandb/core/internal/runwork"
	"github.com/wandb/wandb/core/internal/tensorboard"
	spb "github.com/wandb/wandb/core/pkg/service_go_proto"
)

// RunExitWork implements a special Schedule function for the Exit record.
//
// Otherwise, it defers to the standard WorkRecord implementation.
// It's not yet possible to fully factor out the Exit record's implementation
// from the Handler and Sender.
type RunExitWork struct {
	runwork.Work
	TBHandler *tensorboard.TBHandler
}

type RunExitWorkParams struct {
	Record    *spb.Record
	TBHandler *tensorboard.TBHandler
}

func NewRunExitWork(params RunExitWorkParams) runwork.Work {
	return &RunExitWork{
		Work:      runwork.WorkFromRecord(params.Record),
		TBHandler: params.TBHandler,
	}
}

// Schedule implements Work.Schedule.
//
// We treat records generated by the TB integration as if the client
// sent them. If the integration were in the client, the client would
// stop it and wait for it to finish before sending the Exit record,
// so that it doesn't send any records after Exit.
// This method simulates that.
func (w *RunExitWork) Schedule(wg *sync.WaitGroup, proceed func()) {
	wg.Add(1)
	go func() {
		defer wg.Done()

		w.TBHandler.Finish()
		proceed()
	}()
}
